"""
资源管理器 - 用于管理Pipeline任务的资源分配和限制
"""
import threading
import psutil
import os
from typing import Dict, Any
from core.logger import get_logger

# 获取日志记录器
logger = get_logger(__name__)


class ResourceManager:
    """资源管理器"""

    def __init__(self, global_config: Dict[str, Any] = None):
        """
        初始化资源管理器

        Args:
            global_config: 全局配置
        """
        self.global_config = global_config or {}
        self.resources = self.global_config.get("resources", {})
        self.default_task_resources = self.resources.get("default_task_resources", {
            "cpu_limit": 1.0,
            "memory_limit": "1G",
            "thread_limit": 2
        })
        
        # 资源使用跟踪
        self.resource_usage = {}
        self.lock = threading.Lock()

    def get_node_resources(self, node_id: str, node_config: Dict[str, Any]) -> Dict[str, Any]:
        """
        获取节点的资源配置

        Args:
            node_id: 节点ID
            node_config: 节点配置

        Returns:
            节点资源配置
        """
        # 从节点配置获取资源设置
        node_resources = node_config.get("resources", {})
        
        # 合并默认资源配置和节点特定配置
        final_resources = self.default_task_resources.copy()
        final_resources.update(node_resources)
        
        logger.debug(f"Node {node_id} resources: {final_resources}")
        return final_resources

    def apply_resource_limits(self, node_id: str, resources: Dict[str, Any]):
        """
        应用资源限制到节点

        Args:
            node_id: 节点ID
            resources: 资源配置
        """
        # 获取当前进程
        process = psutil.Process(os.getpid())
        
        # 设置CPU限制（affinity）
        try:
            cpu_limit = resources.get('cpu_limit', 1.0)
            if cpu_limit < 1.0:
                # 对于小于1的核心数，我们限制CPU使用率
                # 注意：psutil不直接支持CPU份额，这里仅作示例
                logger.info(f"Setting CPU limit for node {node_id}: {cpu_limit} cores")
            else:
                # 设置CPU亲和性（限制可用的CPU核心数）
                available_cpus = list(range(int(cpu_limit)))
                if available_cpus:
                    process.cpu_affinity(available_cpus)
                    logger.info(f"Set CPU affinity for node {node_id} to cores: {available_cpus}")
        except Exception as e:
            logger.warning(f"Failed to set CPU limits for node {node_id}: {e}")

        # 记录内存限制信息（psutil不直接支持内存限制，需要使用其他机制）
        memory_limit = resources.get('memory_limit', '1G')
        logger.info(f"Memory limit for node {node_id}: {memory_limit} (enforcement not implemented)")

        # 线程限制信息
        thread_limit = resources.get('thread_limit', 2)
        logger.info(f"Thread limit for node {node_id}: {thread_limit} (enforcement via plugin implementation)")

    def track_resource_usage(self, node_id: str, usage: Dict[str, Any]):
        """
        跟踪资源使用情况

        Args:
            node_id: 节点ID
            usage: 资源使用情况
        """
        # 获取当前进程资源使用情况
        process = psutil.Process(os.getpid())
        current_usage = {
            "cpu_percent": process.cpu_percent(),
            "memory_info": process.memory_info()._asdict(),
            "num_threads": process.num_threads()
        }
        
        # 合并传入的使用情况和实际使用情况
        tracked_usage = usage.copy()
        tracked_usage.update(current_usage)
        
        with self.lock:
            self.resource_usage[node_id] = tracked_usage
            logger.debug(f"Resource usage for node {node_id}: {tracked_usage}")

    def get_resource_usage_report(self) -> Dict[str, Any]:
        """
        获取资源使用报告

        Returns:
            资源使用报告
        """
        with self.lock:
            return self.resource_usage.copy()